{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import sisl\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import tqdm\n",
    "from functools import partial\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the key ideas behind `sisl` was the interaction of a DFT Hamiltonian and the user.  \n",
    "In this example, we will highlight a unique implementation in TBtrans which enables ***any*** kind of user intervention.\n",
    "\n",
    "The idea is a transformation of the Green's function calculation from:\n",
    "\\begin{equation}\n",
    "   \\mathbf G^{-1}(E) = \\mathbf S (E + i\\eta) - \\mathbf H - \\sum_i\\boldsymbol\\Sigma_i\n",
    "\\end{equation}\n",
    "to\n",
    "\\begin{equation}\n",
    "   \\mathbf G^{-1}(E) = \\mathbf S (E + i\\eta) - \\mathbf H - \\delta\\mathbf H - \\sum_i\\boldsymbol\\Sigma_i - \\delta\\boldsymbol \\Sigma\n",
    "\\end{equation}\n",
    "Where $\\mathbf G$ is the Green's function of the system, $\\mathbf S$ is the corresponding overlap matrix, *E* is the energy, $\\mathbf H$ is the Hamiltonian matrix and $\\boldsymbol\\Sigma$ the self-energy/ies. The $\\delta\\mathbf H$ and $\\delta\\boldsymbol\\Sigma$ terms are the perturbation to the Hamiltonian and self-energy/ies, respectively, and can be of any type, i.e. complex and/or real.  \n",
    "The only (important!) difference between $\\delta\\mathbf H$ and $\\delta\\boldsymbol \\Sigma$ is that the former enters the calculation of bond-transmissions, while the latter does not.\n",
    "\n",
    "Since TBtrans by itself does not allow complex Hamiltonians, the above is a way to remove this restriction. One feature this may be used for is to simulate the application of magnetic fields.\n",
    "\n",
    "In the following we will use Peierls substitution on a square tight-binding model:\n",
    "\\begin{equation}\n",
    "   \\mathbf H(\\Phi) = \\mathbf H(0)e^{i \\Phi \\delta x^- \\cdot \\delta y^+ / 2},\n",
    "\\end{equation}\n",
    "where $\\Phi$ is the magnetic flux (in proper units), $\\delta x^- = x_j - x_i$ and $\\delta y^+=y_j + y_i$, with $i$ and $j$ being atomic indices. If you are interested in this substitution, please search literature.\n",
    "\n",
    "---\n",
    "\n",
    "First, create a square lattice and define the on-site and nearest-neighbour couplings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "square = sisl.Geometry([0,0,0], sisl.Atom(1, R=1.0), lattice=sisl.Lattice(1, nsc=[3, 3, 1]))\n",
    "on, nn = 4, -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "H_minimal = sisl.Hamiltonian(square)\n",
    "H_minimal.construct([[0.1, 1.1], [on, nn]])\n",
    "H_elec = H_minimal.tile(100, 1).tile(2, 0)\n",
    "H_elec.set_nsc([3, 1, 1])\n",
    "H_elec.write('ELEC.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "H = H_elec.tile(50, 0)\n",
    "\n",
    "# Make a constriction\n",
    "geom = H.geometry.translate( -H.geometry.center(what='xyz') )\n",
    "# This constriction is based on an example in the Kwant project (called qhe). We however make a slight modification.\n",
    "# Remove some atoms, this will create a constriction of 100 - 40 * 2 = 20 Ang with a Gaussian edge profile\n",
    "remove = (np.abs(geom.xyz[:, 1]) > 50 - 37.5 * np.exp( -(geom.xyz[:, 0] / 12) **2 )).nonzero()[0]\n",
    "# To reduce computations we find the atoms in the constriction such that we can \n",
    "# limit the calculation region.\n",
    "device = (np.abs(geom.remove(remove).xyz[:, 0]) < .6).nonzero()[0]\n",
    "geom.remove(remove).write('test.xyz')\n",
    "\n",
    "# Pretty-print the range of atoms within such smaller device region\n",
    "# Note this is Fortran indices (as required for TBtrans) \n",
    "print(sisl.utils.list2str(device + 1))\n",
    "\n",
    "H = H.remove(remove)\n",
    "H.finalize()\n",
    "H.write('DEVICE.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above printed list of atoms should be specified in the `RUN.fdf` file via the `TBT.Atoms.Device` flag. This is important when one is *only* interested in the transmission, and does not care about the density-of-states. You are always encouraged to select the minimal device region to 1) speed-up computations and 2) drastically reduce memory requirements.  \n",
    "In this regard you should read-in the TBtrans manual about the additional flag `TBT.Atoms.Device.Connect` (you may try, as an additional exercise, to set this flag to `True` and check the resulting difference with the previous calculation).\n",
    "\n",
    "Now we have $\\mathbf H(0)$ with *no* phases due to magnetic fields. As the magnetic field is changing the Hamiltonian, and thus enters the bond-transmission calculations, we have to use the $\\delta\\mathbf H$ term (and *not* $\\delta\\boldsymbol\\Sigma$).\n",
    "\n",
    "The first thing we need to calculate is $\\delta x^- \\cdot\\delta y^+$.\n",
    "Since we already have the Hamiltonian, we can utilize the connections by looping the coupling elements (in a sparse matrix/graph, these are called *edges*). This is *much* cheaper than trying to figure out which atoms are neighbouring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = H.geometry\n",
    "xy = sisl.Hamiltonian(device)\n",
    "for ia in device:\n",
    "    # Get all connecting elements (including it-self)\n",
    "    edges = H.edges(ia)\n",
    "    # Calculate the vector between edges and ia:\n",
    "    #    xyz[edges, :] - xyz[ia, :]\n",
    "    Rij = device.Rij(ia, edges)\n",
    "    # Now calculate the product:\n",
    "    #    (xj - xi) * (yj + yi)\n",
    "    # since we have xj - xi we should get 0 for on-site, as expected.\n",
    "    # Notice that we correct to +yi by adding it twice\n",
    "    xy[ia, edges] = Rij[:, 0] * (Rij[:, 1] + 2 * device.xyz[ia, 1])\n",
    "xy.finalize() # this is only because it will speed up the following calculations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have the coupling dependent phase factor, $\\delta x^-\\cdot\\delta y^+$, and all we need to calculate is $\\delta\\mathbf H$ that transforms $\\mathbf H(0) \\to \\mathbf H(\\Phi)$.  \n",
    "This is easily done as the `Hamiltonian` allows basic element wise operations, i.e. `+`, `-`, `*`, `/` and `**` (the power function).  \n",
    "Your task is to insert the correct mathematical equation below, such that `dH` contains $\\delta \\mathbf H$ for $\\mathbf H(\\Phi) = \\mathbf H(0) + \\delta\\mathbf H$.\n",
    "To help you I have inserted the exponential function. To finalize the equation, you need three terms: `H`, `xy` and `phi` (note we use the reciprocal values as loop variables for ease):\n",
    "<!-- dH = np.exp(-0.5j * phi * xy) * H - H-->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rec_phis = np.arange(1, 51, 4)\n",
    "for i, rec_phi in enumerate(tqdm(rec_phis, unit=\"M\")):\n",
    "    phi = 1 / rec_phi\n",
    "    # Calculate H(Phi)\n",
    "    # TODO: insert the correct mathematical equation below to calculate the correct phase\n",
    "    #dH = np.exp ...\n",
    "    with sisl.get_sile('M_{}.dH.nc'.format(rec_phi), mode='w') as fh:\n",
    "        fh.write_delta(dH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "- Calculate all physical quantities for all different applied magnetic fields.  \n",
    "   Before running the calculations, search the manual on how to save the self-energies (*HINT*: `out-of-core`). By default, TBtrans calculates the self-energies as they are needed. However, if one has the same electrodes, same $k$-grid and same $E$-points for several different runs (as in this case) one can with benefit calculate the self-energies *once*, and then reuse them in subsequent calculations.\n",
    "\n",
    "   To ease the calculation of all magnetic fields, a script (`run.sh`) is located in this directory which will loop over the different fields. Carefully read it to infer which option specifies the $\\delta \\mathbf H$ term.\n",
    "   \n",
    "   Since this example has 14 different setups, each with 51 energy points, it will take some time. Around 30 seconds for the first (includes self-energy calculation), and around 10 seconds for all subsequent setups. So be patient. :)\n",
    "- Secondly, read in all output into the workbook in a list.\n",
    "- **TIME (*HARD*)**:\n",
    "    Choose an energy-point for $B=0$ and $B>0$ such that you have a finite transmission ($T>0$). Next, extract the bond-transmissions and calculate the transmission using the law of current conservation (Kirchhoff's 1st law: $T_{\\mathrm{in}} \\equiv T_{\\mathrm{out}}$). *HINT*: select a line of atoms and calculate the bond-transmissions flowing between this line and its neighbouring line of atoms. Assert that the transmission is the same using both methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create short-hand function\n",
    "gs = sisl.get_sile\n",
    "# No magnetic field\n",
    "tbt0 = gs('siesta.TBT.nc')\n",
    "# All magnetic fields in increasing order\n",
    "tbts = [gs('M_{}/siesta.TBT.nc'.format(rec_phi)) for rec_phi in rec_phis]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Plot the transmission function for all applied fields in the full energy range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do trick with easy plotting utility\n",
    "E = tbt0.E\n",
    "Eplot = partial(plt.plot, E)\n",
    "\n",
    "for rec_phi, tbt in zip(rec_phis, tbts):\n",
    "    Eplot(tbt.transmission(), '--', label='{}'.format(rec_phi));\n",
    "Eplot(tbt0.transmission(), 'k');\n",
    "plt.xlim([E.min(), E.max()]); plt.ylim([0, None]);\n",
    "plt.xlabel('Energy [eV]'); plt.ylabel('Transmisson'); plt.legend(bbox_to_anchor=(1, 1), loc=2);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Make a contour plot of the transmission vs. $\\Phi$ and $E$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = np.stack([tbt.transmission() for tbt in tbts])\n",
    "T[T < 1e-9] = 0 # small numbers are difficult to interpolate (they tend to blow up), so remove them\n",
    "plt.contourf(rec_phis, E, T.T, 20); plt.colorbar(label='Transmission');\n",
    "plt.ylabel('Energy [eV]'); plt.xlabel(r'$1/\\Phi$');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **TIME**: Choose a given magnetic field and create a set of constrictions of different width.  Plot $T(E, \\Phi)$ for increasing widths, fixing $\\Phi$ at a fairly large value (copy codes in `In [4-7]` and adapt). You may decide the Gaussian profile and change it if you want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learned lessons\n",
    "\n",
    "- Advanced construction of geometries by removing subsets of atoms (both for Hamiltonian and geometries)\n",
    "- Creation of $\\delta\\mathbf H$ terms. Note that *exactly* the same method is used for the $\\delta\\boldsymbol\\Sigma$ terms, the only difference is how it is specified in the *fdf*-file (`TBT.dH` vs. `TBT.dSE`). Please look in the TBtrans manual and figure out what the difference is between the two flags?\n",
    "- Supplying *fdf*-flags to TBtrans on the command line to override flags in the input files.\n",
    "- Inform TBtrans to store the self-energies on disk to re-use them in later calculations.\n",
    "- Inform TBtrans to save all output into a sub-folder (and if it does not exist, create it)\n",
    "- Adding complex-valued perturbations to the Hamiltonian matrix. We have not uncovered everything, as the $\\delta$ files may contain $k$-resolved and/or $E$-resolved $\\delta$-terms for full control, but the principles are the same. Search the documentation for `deltancSileTBtrans`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
